import os

import torch
import cv2
from torchvision import transforms

from yolo_resnet34 import YOLOv1_resnet


if __name__ == '__main__':

    net = YOLOv1_resnet()
    # net.load_state_dict(torch.load('YOLOV12021-08-25-22_19_39.pth'))
    # net.load_state_dict(torch.load('YOLOV12021-08-25-22_19_39.pth'))
    net.load_state_dict(torch.load('YOLOV12021-08-25-22_54_05.pth'))

    img = cv2.imread(r'C:\Programs\workspace\deep_learning\data\banana-detection\bananas_val\images\94.png')
    img_initial_shape = img.shape

    target_size = 448

    aug = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize(target_size)
    ])
    img = aug(img)

    img = img.unsqueeze(0)

    net.eval()
    pred = net(img)

    print(pred)
    print(pred[0, 10, :, :])
    print(pred[0, 10, :, :].max())
    print(pred)

    # for x_index in range(7):  # y方向网格循环
    #     for y_index in range(7):  # x方向网格循环
    #         if labels[batch_index, 4, x_index, y_index] == 1:  # 如果包含物体
    #             # 将数据(px,py,w,h)转换为(x1,y1,x2,y2)
    #             # 先将px,py转换为cx,cy，即相对网格的位置转换为标准化后实际的bbox中心位置cx,xy
    #             # 然后再利用(cx-w/2,cy-h/2,cx+w/2,cy+h/2)转换为xyxy形式，用于计算iou
    #
    #             # 预测框中心点在区域中的相对位置 => 预测框左上角在整个图像中的相对位置
    #             bbox1_pred_xyxy = ((pred[batch_index, 0, x_index, y_index] + y_index) / num_gridx
    #                                - pred[batch_index, 2, x_index, y_index] / 2,
    #                                (pred[batch_index, 1, x_index, y_index] + x_index) / num_gridy
    #                                - pred[batch_index, 3, x_index, y_index] / 2,
    #                                (pred[batch_index, 0, x_index, y_index] + y_index) / num_gridx
    #                                + pred[batch_index, 2, x_index, y_index] / 2,
    #                                (pred[batch_index, 1, x_index, y_index] + x_index) / num_gridy
    #                                + pred[batch_index, 3, x_index, y_index] / 2)
    #             bbox2_pred_xyxy = ((pred[batch_index, 5, x_index, y_index] + y_index) / num_gridx
    #                                - pred[batch_index, 7, x_index, y_index] / 2,
    #                                (pred[batch_index, 6, x_index, y_index] + x_index) / num_gridy
    #                                - pred[batch_index, 8, x_index, y_index] / 2,
    #                                (pred[batch_index, 5, x_index, y_index] + y_index) / num_gridx
    #                                + pred[batch_index, 7, x_index, y_index] / 2,
    #                                (pred[batch_index, 6, x_index, y_index] + x_index) / num_gridy
    #                                + pred[batch_index, 8, x_index, y_index] / 2)